package net.maxupeng.asm74;

import net.maxupeng.asm74.antlr.ASM74Parser;
import net.maxupeng.asm74.segment.BSSSegment;
import net.maxupeng.asm74.segment.DataSegment;
import net.maxupeng.asm74.segment.ProcSegment;
import net.maxupeng.asm74.segment.RODataSegment;
import net.maxupeng.asm74.util.AntlrUtil;
import net.maxupeng.asm74.util.LogisimUtil;
import net.maxupeng.asm74.util.Pair;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.*;

import static net.maxupeng.asm74.Compiler.*;

/*
地址空间分配
* ROM
    * $0000-$3FFF   16K
        * init
            * init sp
            * load data default value
                * longjmp
            * jmp main
        * rodata
        * data default value
    * $4000-$7FFF   16K
        * rom page
* RAM
    * $8000-$80FF   编译器占用
    * $8100-$9FFF   8K stack
    * $A000-$BFFF   8K data
    * $C000-$FFEF   16K heap
    * $FFF0-$FFFE   IO
    * $FFFF         页寄存器
 */

// ROM页区域，根据 PAGE 的值换页
public class Main {
    private static final int ROM_PAGE_START = 0x4000;
    private static final int ROM_PAGE_END = 0x7FFF;
    private static final int ROM_PAGE_SIZE = ROM_PAGE_END - ROM_PAGE_START + 1;

    // ROM固定区域，包括初始化代码，RODATA， DATA初始值
    private static final int ROM_FIX_START = 0x0000;
    private static final int ROM_FIX_END = 0x3FFF;
    private static final int ROM_BASE_SIZE = ROM_FIX_END - ROM_FIX_START + 1;

    private static final int RAM_STACK_BASE = 0x9FFF;
    private static final int RAM_STACK_LIMIT = 0x8100;

    private static final int RAM_DATA_START = 0xA000;
    private static final int RAM_DATA_END = 0xBFFF;
    private static final int RAM_DATA_SIZE = RAM_DATA_END - RAM_DATA_START + 1;

    private static final int RAM_HEAP_START = 0xC000;
    private static final int RAM_HEAP_END = 0xFFEF;
    private static final int RAM_HEAP_SIZE = RAM_HEAP_END - RAM_HEAP_START + 1;

    private static final MCodeDefModel model;

    static {
        try {
            model = MCodeDefModel.fromJson(
                    new String(Files.readAllBytes(Paths.get("../mcode.json"))));
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    private static Pair<Integer, Integer> rom_rodata = new Pair<>(0, 0);
    private static Pair<Integer, Integer> ram_rodata = new Pair<>(0, 0);
    private static Pair<Integer, Integer> rom_data = new Pair<>(0, 0);
    private static Pair<Integer, Integer> ram_data = new Pair<>(0, 0);
    private static Pair<Integer, Integer> rom_code = new Pair<>(0, 0);
    private static Pair<Integer, Integer> ram_bss = new Pair<>(0, 0);

    private static OutputStream fos = null;


    private static byte[] asm_proc_line(ProcSegment proc, ASM74Parser.LineContext lineCtx, int addr) {
        if (lineCtx.cmd() != null) {
            if (lineCtx.cmd().getStart().getText().equals("debug")) {
                String s = String.format("[%06x]%02x:%04x", proc.romPage * 32 * 1024 + addr, proc.romPage, addr);
                s = String.format("%1$-" + 55 + "s", s);
                s += " DEBUG: " + lineCtx.cmd().ID();
                if (lineCtx.COMMENT() != null) {
                    s = String.format("%1$-" + 80 + "s", s);
                    s += lineCtx.COMMENT().getText();
                }
                output_list(s);
                return new byte[0];
            }
        }

        String opcode = lineCtx.opcode().getText();
        ASM74Parser.AmContext amCtx = lineCtx.am();
        String amName = AddressMode.getAmName(lineCtx.am());
        byte[] bytes;
        switch (amName) {
            case "default":
                bytes = model.getBinaryCode(opcode, "default", 0, 0);
                break;
            case "imm":
                bytes = model.getBinaryCode(opcode, "imm", 0, evalExpr(amCtx.am_imm().num().expr()));
                break;
            case "direct":
                bytes = model.getBinaryCode(opcode, "direct", evalExpr(amCtx.am_direct().address().expr()), 0);
                break;
            case "indirect":
                bytes = model.getBinaryCode(opcode, "indirect", evalExpr(amCtx.am_indirect().address().expr()), 0);
                break;
            case "indirect_index": {
                bytes = model.getBinaryCode(
                        opcode,
                        "indirect_index",
                        evalExpr(amCtx.am_indirect_index().address().expr()),
                        evalExpr(amCtx.am_indirect_index().num().expr())
                );
                break;
            }
            case "direct_index":
                bytes = model.getBinaryCode(
                        opcode,
                        "direct",
                        evalExpr(amCtx.am_direct_index().address().expr()) + evalExpr(amCtx.am_direct_index().num().expr()),
                        0
                );
                break;
            case "index_indirect":
                bytes = model.getBinaryCode(
                        opcode,
                        "indirect",
                        evalExpr(amCtx.am_index_indirect().address().expr()) + evalExpr(amCtx.am_index_indirect().num().expr()),
                        0
                );
                break;
            default:
                System.out.println("unknown addressing mode");
                System.exit(1);
                bytes = new byte[0];
        }

        String s = String.format("[%06x]%02x:%04x", proc.romPage * 32 * 1024 + addr, proc.romPage, addr);
        s = String.format("%1$-" + 20 + "s", s);
        for (int i = 0; i < bytes.length; i++) {
            s += String.format("%02x ", bytes[i]);
        }
        s = String.format("%1$-" + 35 + "s", s);
        s += proc.romLabelRefMap.getOrDefault(addr, new ArrayList<>()).stream().reduce((a, b) -> a + "," + b).orElse("");
        s = String.format("%1$-" + 55 + "s", s);
        s += " " + opcode + " " + amCtx.getText();
        if (lineCtx.COMMENT() != null) {
            s = String.format("%1$-" + 80 + "s", s);
            s += lineCtx.COMMENT().getText();
        }
        output_list(s);

        return bytes;
    }

    private static void asm_proc(ProcSegment proc) {
        proc.rom.clear();

        for (Map.Entry<String, Integer> entry : proc.romLabelOffsetMap.entrySet()) {
            String label = entry.getKey();
            int offset = entry.getValue();
            proc.romLabelRefMap.computeIfAbsent(proc.romAddr + offset, k -> new ArrayList<>()).add(label);
        }

        int addr = proc.romAddr;
        for (ASM74Parser.LineContext lineCtx : proc.lineCtxs) {
            try {
                byte[] bytes = asm_proc_line(proc, lineCtx, addr);
                addr += bytes.length;
                for (byte b : bytes) {
                    proc.rom.add(b & 0xff);
                }
            } catch (Exception e) {
                System.out.println(lineCtx.getText());
                e.printStackTrace();
                System.exit(1);
            }
        }
    }


    private static void output_list(String str, boolean cr) {
        try {
            if (fos != null) {
                fos.write(str.getBytes());
                if (cr) {
                    fos.write('\n');
                }
            }

            System.out.print(str);
            if (cr) {
                System.out.println();
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    private static void output_list(String str) {
        output_list(str, true);
    }

    private static String gen_ljmp_data() {
        String code_ljmp = """
                .segment "CODE"
                .proc longjmp
                lda TMP0,#2
                sta PAGE
                jmp (TMP0)
                
                """;

        ProcSegment proc_ljmp = ProcSegment.parse(List.of(AntlrUtil.parse(code_ljmp))).get(0);
        asm_proc(proc_ljmp);

        StringBuilder hexBuilder = new StringBuilder();
        for (int i = 0; i < proc_ljmp.rom.size(); i++) {
            int value = proc_ljmp.rom.get(i);
            if (i == 0) {
                hexBuilder.append(String.format("0x%02X", value));
            } else {
                hexBuilder.append(String.format(", 0x%02X", value));
            }
        }

        String data_ljmp = """
                .segment "DATA"
                longjmp: .byte """ + hexBuilder.toString() + "\n\n";

        return data_ljmp;
    }

    private static void hexdump(String name, int[] rom, Pair<Integer, Integer> range) {
        output_list("    \t==================== " + name + " ===================");
        output_list("    \t 0  1  2  3  4  5  6  7  8  9  A  B  C  D  E  F");

        int n = 0;
        byte[] bytes = new byte[16];
        for (int addr = range.first; addr < range.second; addr++) {
            if (n == 0) {
                output_list(String.format("%04x\t", addr), false);
            }
            bytes[n % 16] = (byte) (rom[addr] < 0x20 ? '.' : rom[addr]);
            output_list(String.format("%02x ", rom[addr]), false);

            n++;
            if (n == 16) {
                n = 0;
                output_list("\t" + new String(bytes));
            }
        }
        if (n != 0) {
            output_list("".repeat(3 * (16 - n)) + "\t" + new String(Arrays.copyOf(bytes, n)));
        }
    }

    private static void gen_bp(String lst) throws IOException {
        String bp = lst + "_bp.bin";

        int[] rom = new int[2 << 24];
        Arrays.fill(rom, 0);

        List<String> lines = Files.readAllLines(Paths.get(lst));
        for (String line : lines) {
            if (line.contains(";")) {
                int addr = Integer.parseInt(line.substring(1, 7), 16);
                rom[addr] = 1;
            }
        }

        Files.write(Paths.get(bp), LogisimUtil.convertRomToLogisimFormat(rom).getBytes());
    }

    public static void main(String[] args) throws IOException {
        String[] afiles = {"hello.s"};
        List<String> files = new ArrayList<>();
        for (String file : afiles) {
            files.add(new File(file).getAbsolutePath());
        }
        Collections.sort(files);

        String dir = new File(files.get(0)).getParent();
        String dst = files.size() == 1 ?
                files.get(0).replaceFirst("\\.[^.]+$", ".bin") :
                dir + "/out.bin";
        String lst = files.size() == 1 ?
                files.get(0).replaceFirst("\\.[^.]+$", ".alst") :
                dir + "/out.alst";

        fos = new FileOutputStream(lst);

        List<ASM74Parser.ProgContext> progCtxs = new ArrayList<>();
        for (String file : files) {
            String content = new String(Files.readAllBytes(Paths.get(file)));
            progCtxs.add(AntlrUtil.parse(content));
        }

        // _ljmp in memory
        progCtxs.add(0, AntlrUtil.parse(gen_ljmp_data()));
        List<RODataSegment> rodatas = RODataSegment.parse(progCtxs);

        List<DataSegment> datas = DataSegment.parse(progCtxs);
        List<BSSSegment> bsses = BSSSegment.parse(progCtxs);
        List<ProcSegment> procs = ProcSegment.parse(progCtxs);

        // 分页
        int page = 0;
        int length = 0;
        for (ProcSegment proc : procs) {
            if (proc.romSize >= ROM_PAGE_SIZE) {
                System.out.println("函数大小超过" + ROM_PAGE_SIZE + "：" + proc.name);
                System.exit(1);
            }

            if (length + proc.romSize >= ROM_PAGE_SIZE) {
                page++;
                length = 0;
            }

            proc.romPage = page;
            proc.romAddr = ROM_PAGE_START + length;
            length += proc.romSize;

            setLabelPageAddr(proc.progCtx, proc.name, proc.romPage, proc.romAddr);
            for (Map.Entry<String, Integer> entry : proc.romLabelOffsetMap.entrySet()) {
                String label = entry.getKey();
                int offset = entry.getValue();
                setLabelPageAddr(proc.progCtx, label, proc.romPage, proc.romAddr + offset);
            }
        }

        int romAddr = ROM_FIX_START;
        int romAddrMark = romAddr;

        int ramAddr = RAM_DATA_START;
        int ramAddrMark = ramAddr;

        romAddrMark = romAddr;

        // _init
        String code_init = """
                .segment "CODE"
                .proc _init
                lda #L(STACK_BASE)
                sta SP
                lda #H(STACK_BASE)
                sta SP,#1
                
                """;

        ProcSegment proc_init = ProcSegment.parse(List.of(AntlrUtil.parse(code_init))).get(0);
        proc_init.romAddr = romAddr;
        romAddr += proc_init.romSize;

        // _init_rodata
        int rodata_size = 0;
        for (RODataSegment rodata : rodatas) {
            rodata_size += rodata.romSize;

        }

        String code_init_rodata = """
                        .segment "CODE"
                        .proc _init_rodata
                        lda #%d
                        sta TMP0
                        lda #%d
                        sta TMP0,#1
                        lda #L(rodata_rom_addr)
                        sta TMP1
                        lda #H(rodata_rom_addr)
                        sta TMP1,#1
                        lda #L(rodata_ram_addr)
                        sta TMP2
                        lda #H(rodata_ram_addr)
                        sta TMP2,#1
                _init_rodata0:
                        lda TMP0
                        or TMP0,#1
                        jz _init_rodata1
                        lda (TMP1)
                        sta (TMP2)
                        sec
                        lda TMP0
                        sub #1
                        sta TMP0
                        lda TMP0,#1
                        sub #0
                        sta TMP0,#1
                        clc
                        lda TMP1
                        add #1
                        sta TMP1
                        lda TMP1,#1
                        add #0
                        sta TMP1,#1
                        clc
                        lda TMP2
                        add #1
                        sta TMP2
                        lda TMP2,#1
                        add #0
                        sta TMP2,#1
                        jmp _init_rodata0
                _init_rodata1:
                """.formatted(rodata_size & 0xff, (rodata_size >> 8) & 0xff);
        ProcSegment proc_init_rodata = ProcSegment.parse(List.of(AntlrUtil.parse(code_init_rodata))).get(0);
        proc_init_rodata.romAddr = romAddr;
        romAddr += proc_init_rodata.romSize;

        for (Map.Entry<String, Integer> entry : proc_init_rodata.romLabelOffsetMap.entrySet()) {
            String label = entry.getKey();
            int offset = entry.getValue();
            setLabelPageAddr(null, label, 0, proc_init_rodata.romAddr + offset);
        }

        // _init_data
        for (DataSegment data : datas) {
            data.genInitProc();
            data.getInitProc().romAddr = romAddr;
            romAddr += data.getInitProc().romSize;

            ProcSegment proc = data.getInitProc();
            setLabelPageAddr(proc.progCtx, proc.name, proc.romPage, proc.romAddr);
            for (Map.Entry<String, Integer> entry : proc.romLabelOffsetMap.entrySet()) {
                String label = entry.getKey();
                int offset = entry.getValue();
                setLabelPageAddr(proc.progCtx, label, proc.romPage, proc.romAddr + offset);
            }
        }

        // jmp _main
        String code_jmp_main = """
                .segment "CODE"
                .proc _jmp_main
                lda #L(_main)
                sta TMP0
                lda #H(_main)
                sta TMP0,#1
                lda #P(_main)
                sta TMP0,#2
                jmp longjmp
                
                """;
        ProcSegment proc_jmp_main = ProcSegment.parse(List.of(AntlrUtil.parse(code_jmp_main))).get(0);
        proc_jmp_main.romAddr = romAddr;

        romAddr += proc_jmp_main.romSize;

        rom_code = new Pair<>(romAddrMark, romAddr);

        romAddr += 0x10 - (romAddr & 0xf);     // 对齐

        // rodata
        romAddrMark = romAddr;
        ramAddrMark = ramAddr;

        setLabelPageAddr(null, "rodata_rom_addr", 0, romAddr);
        setLabelPageAddr(null, "rodata_ram_addr", 0, ramAddr);
        for (RODataSegment rodata : rodatas) {
            rodata.romAddr = romAddr;
            romAddr += rodata.romSize;

            rodata.ramAddr = ramAddr;
            ramAddr += rodata.romSize;

            for (Map.Entry<String, Integer> entry : rodata.ramLabelOffsetMap.entrySet()) {
                String label = entry.getKey();
                int offset = entry.getValue();
                setLabelPageAddr(rodata.progCtx, label, 0, rodata.ramAddr + offset);
            }
        }
        rom_rodata = new Pair<>(romAddrMark, romAddr);
        ram_rodata = new Pair<>(ramAddrMark, ramAddr);

        romAddr += 0x10 - (romAddr & 0xf);     // 对齐

        // data default value
        ramAddrMark = ramAddr;
        romAddrMark = romAddr;
        for (DataSegment data : datas) {
            data.romAddr = romAddr;
            romAddr += data.romSize;

            for (Map.Entry<String, Integer> entry : data.romLabelOffsetMap.entrySet()) {
                String label = entry.getKey();
                int offset = entry.getValue();
                setLabelPageAddr(data.progCtx, label, data.romPage, data.romAddr + offset);
            }

            data.ramAddr = ramAddr;
            ramAddr += data.ramSize;

            for (Map.Entry<String, Integer> entry : data.ramLabelOffsetMap.entrySet()) {
                String label = entry.getKey();
                int offset = entry.getValue();
                setLabelPageAddr(data.progCtx, label, data.ramPage, data.ramAddr + offset);
            }
        }
        ram_data = new Pair<>(ramAddrMark, ramAddr);
        rom_data = new Pair<>(romAddrMark, romAddr);

        // bss
        ramAddrMark = ramAddr;
        for (BSSSegment bss : bsses) {
            bss.ramAddr = ramAddr;
            ramAddr += bss.ramSize;

            for (Map.Entry<String, Integer> entry : bss.ramLabelOffsetMap.entrySet()) {
                String label = entry.getKey();
                int offset = entry.getValue();
                setLabelPageAddr(bss.progCtx, label, bss.ramPage, bss.ramAddr + offset);
            }
        }
        ram_bss = new Pair<>(ramAddrMark, ramAddr);

        // 到此，所有的 label 地址都确认完毕，开始使用替换原有的占位符
        for (DataSegment data : datas) {
            for (Map.Entry<Integer, String> entry : data.romLabelRefMap.entrySet()) {
                int offset = entry.getKey();
                String label = entry.getValue();
                int addr = getLabelAddr(data.progCtx, label);
                data.rom.set(offset, addr & 0xff);
                data.rom.set(offset + 1, (addr >> 8) & 0xff);
            }
        }

        // 汇编
        asm_proc(proc_init);
        asm_proc(proc_init_rodata);
        for (DataSegment data : datas) {
            asm_proc(data.getInitProc());
        }
        asm_proc(proc_jmp_main);

        for (ProcSegment proc : procs) {
            asm_proc(proc);
        }

        // 生成 ROM
        int[] rom = new int[2 << 24];
        Arrays.fill(rom, 0);

        proc_init.output_rom(rom);
        proc_init_rodata.output_rom(rom);
        proc_jmp_main.output_rom(rom);

        for (DataSegment data : datas) {
            data.output_rom(rom);
        }
        for (RODataSegment rodata : rodatas) {
            rodata.output_rom(rom);
        }
        for (ProcSegment proc : procs) {
            proc.output_rom(rom);
        }

        hexdump("RODATA", rom, rom_rodata);
        hexdump("DATA", rom, rom_data);

        String segmentInfo = """
                
                ROM CODE :[0x""" + Integer.toString(rom_code.first, 16) + ", 0x" + Integer.toString(rom_code.second, 16) + ") " + (rom_code.second - rom_code.first) + """
                ROM RODATA :[0x""" + Integer.toString(rom_rodata.first, 16) + ", 0x" + Integer.toString(rom_rodata.second, 16) + ") " + (rom_rodata.second - rom_rodata.first) + """
                ROM DATA :[0x""" + Integer.toString(rom_data.first, 16) + ", 0x" + Integer.toString(rom_data.second, 16) + ") " + (rom_data.second - rom_data.first) + """
                
                RAM RODATA :[0x""" + Integer.toString(ram_rodata.first, 16) + ", 0x" + Integer.toString(ram_rodata.second, 16) + ") " + (ram_rodata.second - ram_rodata.first) + """
                RAM DATA :[0x""" + Integer.toString(ram_data.first, 16) + ", 0x" + Integer.toString(ram_data.second, 16) + ") " + (ram_data.second - ram_data.first) + """
                RAM BSSSegment :[0x""" + Integer.toString(ram_bss.first, 16) + ", 0x" + Integer.toString(ram_bss.second, 16) + ") " + (ram_bss.second - ram_bss.first) + """
                
                """;
        System.out.println(segmentInfo);

        /*
         * ROM
         * $0000-$3FFF   16K
         * init
         * rodata
         * data init value
         * $4000-$7FFF   16K
         * rom page
         * ram
         * $8000-$80FF   编译器占用
         * $8100-$9FFF   8K stack
         * $A000-$BFFF   8K data
         * jmp _main
         * rodata
         * data
         * bss
         * $C000-$FFEF   16K heap
         * $FFF0-$FFFF   IO
         */
        if (ram_bss.second > RAM_DATA_END) {
            System.out.println("RAM空间不足");
        }

        if (rom_data.second > ROM_FIX_END) {
            System.out.println("ROM空间不足");
        }

        String output = LogisimUtil.convertRomToLogisimFormat(rom);
        Files.write(Paths.get(dst), output.getBytes());

        try {
            if (fos != null) {
                fos.close();
            }
        } catch (IOException e) {
            e.printStackTrace();
        }

        gen_bp(lst);
    }
}